#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
图像检测模块
提供YOLOv8图像对象检测相关功能
"""

from ultralytics import YOLO
import cv2
import os
import glob
import shutil
import urllib.request
import tkinter as tk
from tkinter import filedialog
import random
import string
import time
from urllib.parse import urlparse
import requests
from io import BytesIO

def load_model(models_dir, model_name='yolov8n.pt'):
    """加载YOLOv8模型，如果不存在则下载"""
    model_path = os.path.join(models_dir, model_name)
    
    if not os.path.exists(model_path):
        print(f"模型文件不存在，将自动下载到 {model_path}")
        model = YOLO(model_name)
        # 确保模型文件保存到models目录
        if os.path.exists(model_name) and not os.path.exists(model_path):
            os.makedirs(os.path.dirname(model_path), exist_ok=True)
            shutil.move(model_name, model_path)
    else:
        model = YOLO(model_path)
    
    print(f"成功加载模型: {model}")
    return model

def download_sample_images(data_dir):
    """准备示例图像到data目录"""
    sample_images = {
        'bus.jpg': 'https://ultralytics.com/images/bus.jpg',
        'zidane.jpg': 'https://ultralytics.com/images/zidane.jpg'
    }
    
    # 检查示例图像是否存在
    missing_images = []
    for img_name in sample_images.keys():
        img_path = os.path.join(data_dir, img_name)
        if not os.path.exists(img_path):
            missing_images.append(img_name)
    
    if missing_images:
        print(f"\n警告：以下示例图像不存在于 {data_dir} 目录中：")
        for img_name in missing_images:
            print(f"  - {img_name}")
        print(f"\n您可以通过以下方式添加示例图像：")
        print(f"  1. 手动下载图像并放入 {data_dir} 目录")
        print(f"  2. 将任意JPG/PNG图像复制到 {data_dir} 目录并重命名")
        print(f"  3. 使用自己的图像文件放入 {data_dir} 目录\n")
    else:
        print(f"✓ 示例图像检查完成，所有示例图像已存在于 {data_dir} 目录中")

def generate_random_suffix(length=6):
    """生成随机字符串作为文件名后缀"""
    chars = string.ascii_lowercase + string.digits
    return ''.join(random.choice(chars) for _ in range(length))

def download_image_from_url(url, save_path=None):
    """从URL下载图像到指定路径，如果save_path为None则只下载到临时文件"""
    # 验证URL格式
    if not url or not url.startswith(('http://', 'https://')):
        print(f"✗ 无效的URL格式: {url}")
        print("  URL必须以http://或https://开头")
        return None, None
    
    try:
        # 解析URL获取文件名
        parsed_url = urlparse(url)
        image_name = os.path.basename(parsed_url.path)
        
        # 如果URL没有明确的文件名，使用时间戳生成一个
        if not image_name or '.' not in image_name:
            image_name = f"image_{int(time.time())}.jpg"
        
        # 如果需要保存，就使用指定路径
        if save_path:
            local_path = save_path
        else:
            # 否则使用临时路径
            temp_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'temp')
            os.makedirs(temp_dir, exist_ok=True)
            local_path = os.path.join(temp_dir, image_name)
        
        print(f"正在从URL下载图像: {url}")
        print(f"这可能需要几秒钟时间...")
        
        # 下载图像 - 使用requests库，设置User-Agent绕过一些反爬虫限制
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Referer': 'https://www.google.com'
        }
        
        # 尝试使用requests下载
        try:
            response = requests.get(url, headers=headers, timeout=15)
            if response.status_code == 200:
                # 验证是否为图像内容
                content_type = response.headers.get('Content-Type', '')
                if not content_type.startswith('image/'):
                    print(f"✗ 下载失败: URL不是图像 (Content-Type: {content_type})")
                    return None, None
                
                # 保存图像文件
                with open(local_path, 'wb') as f:
                    f.write(response.content)
                print(f"✓ 图像下载成功: {local_path}")
                return local_path, image_name
            else:
                print(f"✗ 下载失败，HTTP状态码: {response.status_code}")
                if response.status_code == 404:
                    print("  图像链接不存在，请检查URL是否正确")
                elif response.status_code >= 500:
                    print("  服务器错误，请稍后再试")
                else:
                    print("  无法访问此URL，可能是权限问题或链接已失效")
        except requests.exceptions.Timeout:
            print("✗ 下载超时，请检查您的网络连接或稍后再试")
        except requests.exceptions.ConnectionError:
            print("✗ 连接错误，请检查您的网络连接是否正常")
        except requests.exceptions.InvalidURL:
            print("✗ 无效的URL格式")
        except requests.exceptions.RequestException as e:
            print(f"✗ 请求错误: {e}")
            
        # 如果requests方法失败，尝试使用urllib作为后备方案
        print("尝试使用备用下载方法...")
        try:
            # 创建请求，添加User-Agent
            req = urllib.request.Request(url, headers=headers)
            with urllib.request.urlopen(req, timeout=15) as response:
                # 验证是否为图像内容
                content_type = response.headers.get('Content-Type', '')
                if not content_type.startswith('image/'):
                    print(f"✗ 下载失败: URL不是图像 (Content-Type: {content_type})")
                    return None, None
                
                with open(local_path, 'wb') as f:
                    f.write(response.read())
                print(f"✓ 图像下载成功 (使用备用方法): {local_path}")
                return local_path, image_name
        except urllib.error.URLError as e:
            print(f"✗ URL错误: {e.reason}")
        except urllib.error.HTTPError as e:
            print(f"✗ HTTP错误: {e.code} {e.reason}")
        except TimeoutError:
            print("✗ 连接超时")
        except Exception as e:
            print(f"✗ 备用下载失败: {e}")
    
    except Exception as e:
        print(f"✗ 图像下载失败: {e}")
    
    print("\n请尝试以下解决方法:")
    print("1. 检查URL是否正确")
    print("2. 确保URL指向的是图像文件(jpg, png等)")
    print("3. 检查您的网络连接")
    print("4. 尝试使用其他图像URL")
    print("5. 或者使用本地图像文件代替\n")
    
    return None, None

def detect_image(model, image_path, output_dir, conf_threshold=0.25):
    """使用YOLOv8模型检测图像"""
    if not os.path.exists(image_path):
        print(f"错误: 图像 {image_path} 不存在")
        return None
    
    # 获取图像文件名（不含路径）并添加随机后缀
    image_name = os.path.basename(image_path)
    name_parts = os.path.splitext(image_name)
    random_suffix = generate_random_suffix()
    output_name = f"result_{name_parts[0]}_{random_suffix}{name_parts[1]}"
    output_path = os.path.join(output_dir, output_name)
    
    print(f"正在对图像 {image_path} 进行目标检测...")
    results = model(image_path, conf=conf_threshold)
    
    # 输出检测结果
    detection_info = []
    for r in results:
        boxes = r.boxes
        print(f"检测到 {len(boxes)} 个目标")
        
        # 获取带标注的图像
        im_array = r.plot()
        
        # 为每个检测到的目标收集信息
        for i, box in enumerate(boxes):
            cls = int(box.cls[0])
            conf = float(box.conf[0])
            label = model.names[cls]
            x1, y1, x2, y2 = box.xyxy[0].tolist()
            
            info = {
                'id': i + 1,
                'class': label,
                'confidence': conf,
                'box': (x1, y1, x2, y2)
            }
            detection_info.append(info)
            
            print(f"  目标 {i+1}: 类别 = {label}, 置信度 = {conf:.2f}")
            print(f"    边界框: x1={x1:.1f}, y1={y1:.1f}, x2={x2:.1f}, y2={y2:.1f}")
    
    # 保存结果图像
    cv2.imwrite(output_path, im_array)
    print(f"结果已保存至: {output_path}")
    
    return {
        'original_image': image_path,
        'result_image': output_path,
        'detections': detection_info
    }

def select_image_from_data_dir(data_dir):
    """让用户从数据目录中选择图像"""
    # 获取data目录中的所有图像文件
    image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.webp']
    image_files = []
    
    for ext in image_extensions:
        image_files.extend(glob.glob(os.path.join(data_dir, ext)))
    
    if not image_files:
        print("数据目录中没有图像文件")
        return None
    
    # 显示所有可用的图像
    print("\n可用的图像文件:")
    for i, img_path in enumerate(image_files):
        print(f"  [{i+1}] {os.path.basename(img_path)}")
    
    # 让用户选择
    while True:
        choice = input("\n请选择图像 [输入数字]，或输入 'q' 返回: ")
        if choice.lower() == 'q':
            return None
        
        try:
            index = int(choice) - 1
            if 0 <= index < len(image_files):
                return image_files[index]
            else:
                print("无效的选择，请重新输入")
        except ValueError:
            print("请输入有效的数字")
    
    return None 